package com.mxx.common.security;

import com.mxx.common.exception.EncryptionException;
import com.mxx.common.security.displacer.Displacer;
import com.mxx.common.utils.CharsetUtil;
import com.mxx.common.utils.Strings;
import lombok.extern.slf4j.Slf4j;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.util.Optional;

/**
 * @author: mxx
 * @date: 2022/3/21
 * @time: 10:10
 * @Description: 描述信息
 */
@Slf4j
public class AesCipher {
    private String key;
    private String iv;
    private byte[] secretKeyBytes;
    private byte[] ivBytes;
    private Displacer displacer = Displacer.BASE64;
    private CipherModel cipherModel;
    private PaddingBuild<AesCipher> paddingBuild;


    public static PaddingBuild<AesCipher> ecb(String key) {
        AesCipher aesCipher = new AesCipher();
        aesCipher.key = key;
        aesCipher.cipherModel = new CipherModel(Constant.CipherType.AES, Constant.CipherBlock.ECB, null);
        aesCipher.paddingBuild = new PaddingBuild<>(aesCipher, aesCipher.cipherModel);
        return aesCipher.paddingBuild;
    }

    public static PaddingBuild<AesCipher> cbc(String key, String iv) {
        AesCipher aesCipher = new AesCipher();
        aesCipher.key = key;
        aesCipher.iv = iv;
        aesCipher.cipherModel = new CipherModel(Constant.CipherType.AES, Constant.CipherBlock.CBC, null);
        aesCipher.paddingBuild = new PaddingBuild<>(aesCipher, aesCipher.cipherModel);
        return aesCipher.paddingBuild;
    }

    public AesCipher refreshDisplacer(Displacer displacer) {
        this.displacer = displacer;
        return this;
    }

    public Optional<String> encryptToStr(Object data) {
        Optional<byte[]> optional = this.encryptToByte(data);
        return optional.map(displacer::encoderToString);
    }

    public Optional<byte[]> encryptToByte(Object data) {
        Optional<byte[]> optional = this.execute(data, Cipher.ENCRYPT_MODE);
        return optional;
    }

    public Optional<String> decryptToStr(Object data) {
        Optional<byte[]> optional = this.decryptToByte(data);
        return optional.map(Strings::toString);
    }

    public Optional<byte[]> decryptToByte(Object data) {
        Optional<byte[]> optional = this.execute(data, Cipher.DECRYPT_MODE);
        return optional;
    }

    private Optional<byte[]> execute(Object data, int cipherType) {
        try {
            byte[] secretKeyBytes = this.getSecretKeyBytes();
            byte[] ivBytes = null;
            switch (cipherModel.getBlock()) {
                case CBC:
                    ivBytes = getIvBytes();
                    break;
                case ECB:
                default:
                    break;
            }
            byte[] dataBytes = null;
            if (data instanceof byte[]) {
                dataBytes = (byte[]) data;
            } else {
                switch (cipherType) {
                    case Cipher.ENCRYPT_MODE:
                        dataBytes = ((String) data).getBytes(CharsetUtil.UTF_8);
                        break;
                    case Cipher.DECRYPT_MODE:
                        dataBytes = displacer.decoder((String) data);
                        break;
                    default:
                        throw new EncryptionException("unsupported cipherType : " + cipherType);
                }
            }
            dataBytes = cipherModel.getPadding().complement(dataBytes);

            SecretKey secretKey = new SecretKeySpec(secretKeyBytes, cipherModel.getType().name());
            Cipher cipher = Cipher.getInstance(cipherModel.getModel());
            if (ivBytes == null) {
                cipher.init(cipherType, secretKey);
            } else {
                IvParameterSpec iv = new IvParameterSpec(ivBytes);
                cipher.init(cipherType, secretKey, iv);
            }
            byte[] encryptBytes = cipher.doFinal(dataBytes);
            return Optional.of(encryptBytes);
        } catch (Exception e) {
            e.printStackTrace();
            log.error("aes execute [{}:{}] exception , msg : {}", cipherType, cipherModel.getModel(), e.getMessage());
            return Optional.empty();
        }
    }

    private byte[] getSecretKeyBytes() {
        if (secretKeyBytes == null) {
            this.secretKeyBytes = displacer.decoder(key);
        }
        return this.secretKeyBytes;
    }

    private byte[] getIvBytes() {
        if (ivBytes == null) {
            if (iv != null) {
                this.ivBytes = displacer.decoder(iv);
            } else {
                throw new EncryptionException("cipher iv is null . current block : " + cipherModel.getBlock());
            }
        }
        return this.ivBytes;
    }
}
